package xyz.scootaloo.kami.cloud.handler.pre

import io.vertx.ext.auth.User
import io.vertx.ext.web.Router
import io.vertx.ext.web.RoutingContext
import io.vertx.kotlin.coroutines.await
import xyz.scootaloo.kami.cloud.handler.coroutineHandler
import xyz.scootaloo.kami.cloud.handler.pre.DigestEncryptHandler.verification
import xyz.scootaloo.kami.cloud.lang.currentTimeMillis
import xyz.scootaloo.kami.cloud.lang.getLogger
import xyz.scootaloo.kami.cloud.model.EntityUser
import xyz.scootaloo.kami.cloud.model.core.UserPrincipal
import xyz.scootaloo.kami.cloud.service.AccountService
import xyz.scootaloo.kami.cloud.util.Convert

/**
 * HTTP摘要认证服务
 *
 * 服务端向客户端发起挑战:
 *
 * 将随机数,域和服务端支持的算法列表发送给客户端;
 * 同时响应状态码设置为401(Unauthorized);
 *
 * 客户端收到401状态码, 接收挑战:
 *
 * 从服务端支持的算法列表中选择一个算法, 计算出密码和其他数据的摘要;
 * 将摘要和所使用的算法发回给服务器, 如果客户端要对服务器进行认证, 可以发送客户端随机数;
 *
 * 服务端对客户端身份进行验证:
 * [verification]
 *
 * example:
 * ```text
 * HTTP /1.1 401 Unauthorized
 * WWW-Authenticate: Digest
 *     realm="<realm-value>"
 *     nonce="<nonce-value>"
 *     [domain="<list-of-URIs>"]
 *     [opaque="<opaque-token-value>"]
 *     [stale=<true-or-false>]
 *     [algorithm=<digest-algorithm>]
 *     [qop="<list-of-qop-values>"]
 *     [<extension-directive>]
 * ```
 *
 * @author flutterdash@qq.com
 * @since 2022/4/17 11:40
 */
object DigestEncryptHandler {

    private const val AUTHENTICATE = "WWW-Authenticate" // 质询
    private const val AUTHORIZATION = "Authorization"   // 响应
    private const val AUTHENTICATION_INFO = "Authentication-Info"
    private const val UNAUTHORIZED = 401

    private const val DIGEST_PREFIX = "Digest"
    private const val C_USER = "user"
    private const val C_USERNAME = "username"
    private const val C_QOP = "qop"
    private const val C_RSP_AUTH = "rspauth"
    private const val C_CLIENT_NONCE = "cnonce"
    private const val C_RESPONSE = "response"
    private const val C_NONCE_COUNTER = "nc"
    private const val C_NONCE = "nonce"
    private const val C_URI = "uri"
    private const val C_REALM = "realm"
    private const val C_STALE = "stale"

    private const val DEF_REALM = "kami-cloud-final"
    private const val DEF_QOP = "auth"
    private const val DEF_PRIVATE_KEY = "fly me to the moon"
    private const val MAX_NONCE_AGE_SECONDS = 60

    private val accountService = AccountService()

    private val log by lazy { getLogger("digest-router") }

    fun handler(router: Router) {
        router.route().coroutineHandler(log) { ctx ->
            val success = smartHandler(ctx)
            if (!success) {
                return@coroutineHandler
            }
            ctx.next()
        }
    }

    private suspend fun smartHandler(ctx: RoutingContext): Boolean {
        val authorization = ctx.request().headers()[AUTHORIZATION]
        return if (authorization == null) {
            challenge(ctx)
        } else {
            verification(ctx, authorization)
        }
    }

    private fun challenge(ctx: RoutingContext, stale: Boolean = false): Boolean {
        val response = ctx.response()
        response.headers().add(AUTHENTICATE, Tools.challenge(stale))
        response.statusCode = UNAUTHORIZED
        ctx.end()
        return false
    }

    private suspend fun verification(ctx: RoutingContext, authorization: String): Boolean {
        val method = ctx.request().method().toString()
        val response = ctx.response()
        val authHeader = Tools.parseAuthHeader(authorization, method)
        var stale = false
        if (authHeader != null) {
            val isNonceValid = Tools.validateNonce(authHeader.nonce)

            if (isNonceValid == null) {
                // 随机数过期
                stale = true
            } else if (isNonceValid) {
                // nonce验证通过, 接下来验证客户端是否真的知道账号的密码(根据算法比较摘要)
                val entityUser = accountService.findUser(authHeader.username).await()
                if (entityUser != null) {
                    // 账号存在, 开始计算摘要
                    val computedResponse = Tools.computedResponse(authHeader, entityUser.password)
                    if (computedResponse == authHeader.response) {
                        response.putHeader(
                            AUTHENTICATION_INFO, Tools.authorizationInfo(authHeader, entityUser.password)
                        )
                        setUserInContext(ctx, entityUser)
                        return true
                    }
                } else {
                    // 账号不存在
                }
            }
        }

        return challenge(ctx, stale)
    }

    // 设置上下文中的用户信息
    private fun setUserInContext(ctx: RoutingContext, entityUser: EntityUser) {
        ctx.setUser(User.create(UserPrincipal.json(entityUser)))
    }

    object Tools {
        fun parseAuthHeader(authorization: String, method: String): AuthorizationHeader? {
            if (authorization.startsWith(DIGEST_PREFIX)) {
                val rest = authorization.substring(DIGEST_PREFIX.length + 1)
                    .replace("\"", "")
                return parseAuthHeaderCore(rest, method)
            }
            return null
        }

        fun challenge(stale: Boolean = false): String {
            val parts = mutableListOf(
                Triple(C_REALM, DEF_REALM, true),
                Triple(C_QOP, DEF_QOP, true),
                Triple(C_NONCE, newNonce(), true)
            )
            if (stale) {
                parts.add(Triple(C_STALE, "true", false))
            }
            return "$DIGEST_PREFIX ${parts.format()}"
        }

        fun authorizationInfo(header: AuthorizationHeader, password: String): String {
            return listOf(
                Triple(C_QOP, DEF_QOP, true),
                Triple(C_RSP_AUTH, rspAuth(header, password), true),
                Triple(C_CLIENT_NONCE, header.clientNonce, true),
                Triple(C_NONCE_COUNTER, header.nonceCounter, false)
            ).format()
        }

        private fun parseAuthHeaderCore(authorization: String, method: String): AuthorizationHeader? {
            val params = HashMap<String, String>()
            for (item in authorization.split(',')) {
                val idx = item.indexOf('=')
                if (idx == -1)
                    continue
                val key = item.substring(0, idx).trim()
                val value = item.substring(idx + 1)
                params[key] = value
            }

            return try {
                AuthorizationHeader(
                    username = params[C_USER] ?: params[C_USERNAME]!!,
                    realm = params[C_REALM]!!,
                    method = method,
                    nonce = params[C_NONCE]!!,
                    uri = params[C_URI]!!,
                    nonceCounter = params[C_NONCE_COUNTER]!!,
                    clientNonce = params[C_CLIENT_NONCE]!!,
                    response = params[C_RESPONSE]!!,
                    qop = params[C_QOP]!!
                )
            } catch (nullErr: NullPointerException) {
                null
            }
        }

        private fun rspAuth(header: AuthorizationHeader, password: String): String {
            val a1Hash = header.run { md5("$username:$realm:$password") }
            val a2Hash = header.run { md5(":$uri") }
            return header.run {
                md5("$a1Hash:$nonce:$nonceCounter:$clientNonce:$qop:$a2Hash")
            }
        }

        fun computedResponse(header: AuthorizationHeader, password: String): String {
            val a1Hash = header.run { md5("$username:$realm:$password") }
            val a2Hash = header.run { md5("$method:$uri") }
            return header.run {
                md5("$a1Hash:$nonce:$nonceCounter:$clientNonce:$qop:$a2Hash")
            }
        }

        // 检查密文是否被篡改
        fun validateNonce(nonce: String): Boolean? = try {
            val plainNonce = Convert.base64decode(nonce).trim('\"')
            val timestamp = plainNonce.substring(0, plainNonce.indexOf(' '))
            if (nonce == newNonce(timestamp)) {
                if (currentTimeMillis() - timestamp.toLong() > (MAX_NONCE_AGE_SECONDS * 1000)) {
                    null
                } else {
                    true
                }
            } else {
                false
            }
        } catch (e: Throwable) {
            false
        }

        // "timestamp md5(timestamp:private_key)"
        private fun newNonce(timestamp: String = currentTimeMillis().toString()): String {
            val secret = md5("$timestamp:$DEF_PRIVATE_KEY")
            return Convert.base64encode("\"$timestamp $secret\"")
        }

        private fun List<Triple<String, String, Boolean>>.format(): String {
            return joinToString(",") {
                if (it.third) {
                    "${it.first}=\"${it.second}\""
                } else {
                    "${it.first}=${it.second}"
                }
            }
        }

        private fun md5(data: String): String {
            return Convert.md5(data)
        }
    }

    data class AuthorizationHeader(
        val username: String,
        val realm: String,
        val method: String,
        val uri: String,
        val nonce: String,
        val nonceCounter: String,
        val clientNonce: String,
        val qop: String,
        val response: String
    )
}